#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2019/4/1 14:18
# @Author  : Seven
# @File    : visualization.py
# @Software: PyCharm
# function
from inference import detection
import cv2
from classify_image import run_inference_on_image
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import os
from visualization_utils import draw_bounding_box_on_image_array
rootdir = os.getcwd()


def darw(image, boxes, name='', cut=True):
    """
    在图片上画出车辆的框和打上车辆的型号
    :param image: 图片原数据
    :param boxes: 车辆的bbox信息
    :param name: 车辆的型号信息
    :param cut: 是否剪切图片
    :return: 返回最终标注好的图片
    """
    # 获取图片的维度
    shape = image.shape
    im_height, im_width, _ = shape
    # 轮循每个车辆进行标注
    for idx, box in enumerate(boxes):
        ymin, xmin, ymax, xmax = box
        (left, right, top, bottom) = (xmin * im_width, xmax * im_width,
                                      ymin * im_height, ymax * im_height)
        if cut:  # 剪切出车辆信息，方便进行车辆分类
            car_image = image[int(top):int(bottom), int(left):int(right)]
            cv2.imwrite(os.path.join(rootdir, "car/%d.jpg" % idx), car_image)
        else:  # 进行车辆标注
            draw_bounding_box_on_image_array(image, ymin, xmin, ymax, xmax, )
            cv2.rectangle(image, (int(left), int(top)), (int(right), int(bottom)), (255, 0, 0), thickness=4)
            image = draw_text(image, (int(left), int(top)), '{}:{}'.format(idx+1, name[idx]))
    return image


def draw_text(img, position, text):
    """
    在车辆的位置标准车辆型号信息
    :param img: 有车辆位置信息的图片
    :param position: 标注位置
    :param text: 标注信息
    :return: 返回标注好的图片
    """
    cv2img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # cv2和PIL中颜色的hex码的储存顺序不同
    pilimg = Image.fromarray(cv2img)
    # PIL图片上打印汉字
    draw = ImageDraw.Draw(pilimg)  # 图片上打印
    font = ImageFont.truetype(os.path.join(rootdir, 'data/SimHei.ttf'), 30, encoding="utf-8")
    draw.text(position, text, font=font, fill='red')
    cv2charimg = cv2.cvtColor(np.array(pilimg), cv2.COLOR_RGB2BGR)
    return cv2charimg


def classifer():
    """
    车辆分类
    :return: 返回车辆的类别信息
    """
    car_path = os.path.join(rootdir, 'car')
    cars = os.listdir(car_path)
    name = []
    score = []
    for car in cars:
        image_path = os.path.join(car_path, car)
        n, scores = run_inference_on_image(image_path)
        name.append(n)
        score.append(scores)
        os.remove(image_path)
    return name, score


def main(image_path):
    image_np = cv2.imread(image_path)
    print('开始检测车辆')
    boxes, scores = detection(image_path)

    if boxes:
        print("检测到的车辆数目：%s" % len(boxes))
        darw(image_np, boxes)
        print("开始车辆型号识别")
        name, class_score = classifer()
        text = []
        for n, s in zip(name, scores):
            text.append("检测到车辆: {}".format(n))
        print("检测到的车辆型号:", text)
        image_car = darw(image_np, boxes, text, cut=False)
        cv2.imwrite(os.path.join(rootdir, 'result.jpg'), image_car)
        # cv2.imshow('end', image_car)
        return False, scores, class_score
    else:
        print("没有检测到车辆！！！")
        return True, 0, 0


if __name__ == '__main__':
    image = 'images/2333.jpg'
    main(image)
    cv2.waitKey(0)
